import { providers } from 'shared/providers'
import { create } from 'zustand'
import { persist } from 'zustand/middleware'
import { immer } from 'zustand/middleware/immer'

export interface ChatMessage {
  id: string
  role: 'user' | 'assistant' | 'system'
  content: string
  reasoning_content?: string
  timestamp: number
  isError?: boolean
}

export type AIProvider = keyof typeof providers | 'custom'

type APIKeys = {
  [key in AIProvider]: string
}

export interface ProviderConfig {
  provider: AIProvider
  model: string
  modelPreferences: {
    [key in AIProvider]: string
  }
}

export interface LiveStreamConfig {
  positivePrompt: string // 正向提示词
  negativePrompt: string // 反向提示词
  role: string // 角色
  productParams: string // 产品参数
}

export interface LiveStreamConfigPreset {
  id: string
  name: string
  config: LiveStreamConfig
}

type Status = 'ready' | 'waiting' | 'replying'

interface AIChatStore {
  messages: ChatMessage[]
  status: Status
  apiKeys: APIKeys
  config: ProviderConfig
  customBaseURL: string
  liveStreamConfigs: LiveStreamConfigPreset[]
  activeLiveStreamConfigId: string | null
  addLiveStreamConfig: (preset: Omit<LiveStreamConfigPreset, 'id'>) => void
  updateLiveStreamConfig: (id: string, preset: Partial<LiveStreamConfigPreset>) => void
  deleteLiveStreamConfig: (id: string) => void
  setActiveLiveStreamConfig: (id: string | null) => void
  getActiveLiveStreamConfig: () => LiveStreamConfig | null
  setCustomBaseURL: (url: string) => void
  setConfig: (config: Partial<ProviderConfig>) => void
  setApiKey: (provider: AIProvider, key: string) => void
  addMessage: (message: Omit<ChatMessage, 'id' | 'timestamp'>) => void
  appendToChat: (chunk: string) => void
  appendToReasoning: (chunk: string) => void
  tryToHandleEmptyMessage: (message: string) => void
  setMessages: (messages: ChatMessage[]) => void
  setStatus: (status: Status) => void
  clearMessages: () => void
  autoScroll: boolean
  setAutoScroll: (value: boolean) => void
}

export const useAIChatStore = create<AIChatStore>()(
  persist(
    immer<AIChatStore>(set => {
      const modelPreferences = Object.keys(providers).reduce(
        (acc, provider) => {
          acc[provider as AIProvider] =
            providers[provider as keyof typeof providers].models[0] || ''
          return acc
        },
        {} as Record<AIProvider, string>,
      )

      const apiKeys = Object.keys(providers).reduce(
        (acc, provider) => {
          acc[provider as AIProvider] = ''
          return acc
        },
        {} as Record<AIProvider, string>,
      )

      return {
        messages: [],
        status: 'ready',
        config: {
          provider: 'deepseek',
          model: providers.deepseek.models[0],
          modelPreferences,
        },
        apiKeys,
        customBaseURL: '',
        liveStreamConfigs: [],
        activeLiveStreamConfigId: null,
        addLiveStreamConfig: preset => {
          set(state => {
            const newPreset: LiveStreamConfigPreset = {
              ...preset,
              id: crypto.randomUUID(),
            }
            state.liveStreamConfigs.push(newPreset)
            // 如果是第一个配置，自动激活
            if (state.liveStreamConfigs.length === 1) {
              state.activeLiveStreamConfigId = newPreset.id
            }
          })
        },
        updateLiveStreamConfig: (id, updates) => {
          set(state => {
            const index = state.liveStreamConfigs.findIndex(p => p.id === id)
            if (index !== -1) {
              state.liveStreamConfigs[index] = {
                ...state.liveStreamConfigs[index],
                ...updates,
                config: {
                  ...state.liveStreamConfigs[index].config,
                  ...(updates.config || {}),
                },
              }
            }
          })
        },
        deleteLiveStreamConfig: id => {
          set(state => {
            state.liveStreamConfigs = state.liveStreamConfigs.filter(
              p => p.id !== id,
            )
            // 如果删除的是当前激活的配置，清空激活状态
            if (state.activeLiveStreamConfigId === id) {
              state.activeLiveStreamConfigId =
                state.liveStreamConfigs[0]?.id || null
            }
          })
        },
        setActiveLiveStreamConfig: id => {
          set(state => {
            state.activeLiveStreamConfigId = id
          })
        },
        getActiveLiveStreamConfig: (): LiveStreamConfig | null => {
          const state: AIChatStore = useAIChatStore.getState()
          if (!state.activeLiveStreamConfigId) return null
          const preset: LiveStreamConfigPreset | undefined = state.liveStreamConfigs.find(
            (p: LiveStreamConfigPreset) => p.id === state.activeLiveStreamConfigId,
          )
          return preset?.config || null
        },
        setCustomBaseURL: (url: string) => {
          set(state => {
            state.customBaseURL = url
          })
        },
        setConfig: (config: Partial<ProviderConfig>) => {
          set(state => {
            if (config.provider) {
              const newModel =
                config.model || state.config.modelPreferences[config.provider]
              state.config.provider = config.provider
              state.config.model = newModel
              state.config.modelPreferences[config.provider] = newModel
            } else if (config.model) {
              state.config.model = config.model
              state.config.modelPreferences[state.config.provider] =
                config.model
            }
          })
        },
        setApiKey: (provider: AIProvider, key: string) => {
          set(state => {
            state.apiKeys[provider] = key
          })
        },
        addMessage: (message: Omit<ChatMessage, 'id' | 'timestamp'>) => {
          set(state => {
            state.messages.push({
              ...message,
              id: crypto.randomUUID(),
              timestamp: Date.now(),
            })
          })
        },
        appendToChat: (chunk: string) => {
          set(state => {
            if (
              state.messages[state.messages.length - 1].role !== 'assistant'
            ) {
              state.messages.push({
                role: 'assistant',
                content: chunk,
                id: crypto.randomUUID(),
                timestamp: Date.now(),
              })
            } else {
              state.messages[state.messages.length - 1].content += chunk
            }
          })
        },
        appendToReasoning: (chunk: string) => {
          set(state => {
            if (
              state.messages[state.messages.length - 1].role !== 'assistant'
            ) {
              state.messages.push({
                role: 'assistant',
                reasoning_content: chunk,
                content: '',
                id: crypto.randomUUID(),
                timestamp: Date.now(),
              })
            } else {
              state.messages[state.messages.length - 1].reasoning_content +=
                chunk
            }
          })
        },
        tryToHandleEmptyMessage: (message: string) => {
          set(state => {
            const lastRole = state.messages[state.messages.length - 1]?.role
            if (!lastRole || lastRole === 'user') {
              state.messages.push({
                role: 'assistant',
                content: message,
                id: crypto.randomUUID(),
                timestamp: Date.now(),
                isError: true,
              })
            }
          })
        },
        setMessages: (messages: ChatMessage[]) => {
          set(state => {
            state.messages = messages
          })
        },
        setStatus: (status: Status) => {
          set(state => {
            state.status = status
          })
        },
        clearMessages: () => {
          set(state => {
            state.messages = []
          })
        },
        autoScroll: true,
        setAutoScroll: (value: boolean) => set({ autoScroll: value }),
      }
    }),
    {
      name: 'ai-chat-storage',
      partialize: state => ({
        apiKeys: state.apiKeys,
        config: state.config,
        customBaseURL: state.customBaseURL,
        liveStreamConfigs: state.liveStreamConfigs,
        activeLiveStreamConfigId: state.activeLiveStreamConfigId,
      }),
    },
  ),
)
